/**
 * Copyright (C) 2024 Intel Corporation
 *
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not
 * use this file except in compliance with the License.  You may obtain a copy
 * of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software distributed
 * under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
 * CONDITIONS OF ANY KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations under the License.
 *
 *
 * SPDX-License-Identifier: Apache-2.0
 */

/* clang-format off */

#include <numeric>
#include <initializer_list>
#include "backends/tofino/bf-p4c/common/asm_output.h"  // canon_name
#include "backends/tofino/bf-p4c/parde/clot/clot_info.h"  // ClotInfo
#include "backends/tofino/bf-p4c/device.h"
#include "resources_clot.h"

namespace BFN {

bool ClotResourcesLogging::usingClots() const {
    return Device::numClots() > 0 && BackendOptions().use_clot;
}

std::vector<ClotResourcesLogging::ClotUsage *> &ClotResourcesLogging::getUsageData(gress_t gress,
                                                                                   unsigned tag) {
    return usageData[gress][tag];
}

bool ClotResourcesLogging::preorder(const IR::BFN::LoweredParserState *state) {
    if (!usingClots()) return false;

    auto parserIR = findContext<IR::BFN::LoweredParser>();
    BUG_CHECK(parserIR, "State does not belong to a parser? %1%", state);

    for (const auto *match : state->transitions) {
        collectClotUsages(match, state, parserIR->gress);
    }

    return true;
}

void ClotResourcesLogging::end_apply() {
    if (usingClots()) {
        logClotUsages();
    }
    collected = true;
}

void ClotResourcesLogging::collectClotUsages(const IR::BFN::LoweredParserMatch *match,
                                             const IR::BFN::LoweredParserState *state,
                                             gress_t gress) {
    for (auto *ck : match->checksums) {
        if (auto *csum = ck->to<IR::BFN::LoweredParserChecksum>())
            if (csum->type == IR::BFN::ChecksumMode::CLOT)
                clotTagToChecksumUnit[csum->clot_dest.tag] = csum->unit_id;
    }

    for (auto *stmt : match->extracts) {
        if (auto *extract = stmt->to<IR::BFN::LoweredExtractClot>()) {
            // Don't generate JSON for the extracted CLOT if the CLOT is spilled.
            if (!extract->is_start) continue;

            collectExtractClotInfo(extract, state, gress);
        }
    }
}

static bool areFieldListsEqual(const std::vector<ClotResourcesLogging::ClotField *> &clot1,
                               const std::vector<ClotResourcesLogging::ClotField *> &clot2) {
    if (clot1.size() != clot2.size()) return false;

    for (std::size_t i = 0; i < clot1.size(); i++) {
        // NOTE: Offset is inferred from lsb and msb, so we don't have to check it
        if (clot1[i]->get_name() != clot2[i]->get_name() &&
            clot1[i]->get_field_lsb() != clot2[i]->get_field_lsb() &&
            clot1[i]->get_field_msb() != clot2[i]->get_field_msb())
            return false;
    }

    return true;
}

static bool areClotUsagesEqual(ClotResourcesLogging::ClotUsage *u1,
                               ClotResourcesLogging::ClotUsage *u2) {
    // NOTE: Parameters cannot be const because get_fields_lists() doesn't
    // have const variant (and it is autogenerated code)
    return u1->get_length() == u2->get_length() && u1->get_offset() == u2->get_offset() &&
           areFieldListsEqual(u1->get_field_lists(), u2->get_field_lists());
}

void ClotResourcesLogging::collectExtractClotInfo(const IR::BFN::LoweredExtractClot *extract,
                                                  const IR::BFN::LoweredParserState *state,
                                                  gress_t gress) {
    const auto tag = extract->dest->tag;
    const bool hasChecksum = clotTagToChecksumUnit.count(tag) > 0;
    const auto *source = extract->source->to<IR::BFN::LoweredPacketRVal>();
    const auto bytes = source->range;
    const auto currentExtractLength = bytes.size();
    const auto offset = bytes.lo;
    const std::string issueState = state->name.c_str();
    const Clot *clot = clotInfo.parser_state_to_clot(state, tag);
    BUG_CHECK(clot, "No fields extracted to this tag %1%", tag);

    cstring state_name = clotInfo.sanitize_state_name(extract->higher_parser_state->name, gress);
    BUG_CHECK(clot->parser_state_to_slices().count(state_name),
              "Tried to get field slices in %1% from a parser state it is not present in (%2%)",
              *clot, state_name);
    const std::vector<const PHV::FieldSlice *> slices =
        clot->parser_state_to_slices().at(state_name);

    // Aggregate lengths of extracted fields in bytes
    const auto length = std::accumulate(slices.begin(), slices.end(), 0,
                                        [](int r, const PHV::FieldSlice *const slice) -> int {
                                            return r + slice->size();
                                        }) /
                        8;

    const bool spilledExtraction = length > currentExtractLength;

    if (spilledExtraction) {
        LOG6("Spilled CLOT extraction detected for " << gress << "::CLOT" << tag);
        LOG6("  State extracts: " << currentExtractLength << "B, total length: " << length << "B");
        LOG7("  Fields extracted:");
        for (auto *s : slices) {
            LOG7("    " << *s);
        }
        // We might want to log extra info about spilled CLOTs in the future
    }

    /*
     Loop over all usages for this particular CLOT tag
     All properties are the same then another state is issuing
     the same data to the same CLOT tag. So add its name to array
     of state names.
     If no match is found then current state is issuing different,
     mutually exclusive data to the same CLOT tag, so create new entry
     for it.
     */
    auto newUsage = logExtractClotInfo(state_name, hasChecksum, length, offset, tag, clot);
    auto &usages = getUsageData(gress, tag);

    for (auto &usage : usages) {
        if (areClotUsagesEqual(usage, newUsage)) {
            usage->append_issue_states(issueState);
            return;
        }
    }

    usages.push_back(newUsage);
    usages.back()->append_issue_states(issueState);
    LOG6("Adding new ClotUsage for " << gress << "::CLOT" << tag);
}

void ClotResourcesLogging::logClotUsages() {
    for (auto gress : {INGRESS, EGRESS}) {
        for (auto &kv : usageData[gress]) {
            for (auto &usageData : kv.second) {
                clotUsages[gress]->append_clots(usageData);
            }
        }
    }
}

ClotResourcesLogging::ClotUsage *ClotResourcesLogging::logExtractClotInfo(cstring parser_state,
                                                                          bool hasChecksum,
                                                                          int length, int offset,
                                                                          unsigned tag,
                                                                          const Clot *clot) {
    auto usage = new ClotUsage(hasChecksum, length, offset, tag);

    for (auto *slice : clot->parser_state_to_slices().at(parser_state)) {
        const std::string name = cstring(canon_name(slice->field()->name)).c_str();
        const auto msb = slice->range().hi;
        const auto lsb = slice->range().lo;
        const auto clot_offset = clot->bit_offset(parser_state, slice);
        auto clot_field = new ClotField(clot_offset, lsb, msb, name);
        usage->append_field_lists(clot_field);
    }

    return usage;
}

std::vector<ClotResourcesLogging::ClotResourceUsage *> ClotResourcesLogging::getLoggers() {
    BUG_CHECK(collected, "Trying to get clot log without applying inspector to pipe node.");
    return clotUsages;
}

ClotResourcesLogging::ClotResourcesLogging(const ClotInfo &clotInfo) : clotInfo(clotInfo) {
    using ClotEligibleField = Resources_Schema_Logger::ClotEligibleField;

    if (!usingClots()) return;

    // Once Multi Parser support is added in T2NA, clot support will have
    // ClotInfo per parser and the code below should change to use a vector
    // of clots per parser similar to P4iParser. Currently this assumes a
    // single parser scenario
    // Initialize clot structures for ingress & egress
    for (auto gress : {INGRESS, EGRESS}) {
        clotUsages[gress] = new ClotResourceUsage(::toString(gress).c_str(), Device::numClots());
    }
    usageData.resize(2);

    // Populate information for CLOT-eligible fields.
    for (auto field : *clotInfo.clot_eligible_fields()) {
        const auto bitWidth = field->size;
        const bool isReadonly = clotInfo.is_readonly(field);
        const bool isModified = clotInfo.is_modified(field);
        const bool isChecksum = clotInfo.is_checksum(field);
        const std::string name = cstring(canon_name(field->name)).c_str();

        int numBitsInClots = 0;
        for (auto &kv : *clotInfo.slice_clots(field)) {
            numBitsInClots += kv.first->size();
        }

        int numBitsInPhvs = 0;
        field->foreach_alloc([&](const PHV::AllocSlice &alloc) { numBitsInPhvs += alloc.width(); });

        auto cef = new ClotEligibleField(bitWidth, isChecksum, isModified, isReadonly, name,
                                         numBitsInClots, numBitsInPhvs);
        clotUsages[field->gress]->append_clot_eligible_fields(cef);
    }
}

}  // namespace BFN

/* clang-format on */
